args_ok = ""

bat_content = '''@echo off
color F0

set PYTHON=
set GIT=
set VENV_DIR=
set COMMANDLINE_ARGS=<args>

call webui.bat
'''

print('''
===========================
   启动脚本生成器
===========================
请先选择你的显存大小

1) 2G显存
2) 4G显存
3) 6G显存
4) 8G显存及以上
''')

vram_choice = input("请输入上面的数字后按回车(1-4): ")

if vram_choice == "1":
    args_ok += "--lowvram "
elif vram_choice == "2":
    args_ok += "-- medvram "
elif vram_choice == "3":
    args_ok += "--medvram "
elif vram_choice == "4":
    pass
else:
    print("选择的不对 请重试!")
    exit(0)

print('''==========================
请选择你需要的启动参数

1) 解决16系显卡生成图片黑 (单精度)
2) xformers (可能会使占用变低)
3) deepdanbooru (识别tag用, 训练可能会用到)
4) 暗色主题
5) 自动打开网页
6) OutPaint(√api)

提示: 如果你想同时选择1和2只需要输入12就可以了
''')

args_choice = input("请输入上面的数字后按回车: ")

arg_dic = {
    "1": "--precision full --no-half ",
    "2": "--xformers ",
    "3": "--deepdanbooru ",
    "4": "--theme dark ",
    "5": "--autolaunch ",
    "6": "--api --cors-allow-origins=https://www.painthua.com "
}

for c in args_choice:
    if c in arg_dic:
        args_ok += arg_dic[c]
    else:
        print(f"无法识别这个参数: {c}")

print(f"最终使用的启动参数: {args_ok}")

final = bat_content.replace("<args>", args_ok)

with open("start_webui启动.bat", "w", encoding="utf-8") as f:
    f.write(final)

print("生成完成~ 请使用start_webui启动.bat启动")